# -*-Python-*-

# t5.1.n4w10.base, t5.1.n4w10.large, and t5.1.n4w10.xl are variations on the
# t5.1.1 models.  Each of the encoder and decoder consists of 14 layer groups,
# with the last ten twice as "wide" as the first four.  (double d_ff and
# num_heads). Parameter count and computation are kept similar to the
# corresponding t5.1.1 models.  For the base model, this increases the number of
# layers, resulting in better quality, and for the large and xl models, this
# decreases the number of layers from 24 to 14, decreasing quality, but also
# decreasing the amount of communication necessary for model parallelism.

include 'models/t5.1.n4w10.base.gin'

d_model = 1024
d_ff = 5120
num_heads = 32
utils.tpu_mesh_shape.model_parallelism = 2

encoder/transformer.make_layer_stack.layers = [
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
]
decoder/transformer.make_layer_stack.layers = [
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.EncDecAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.EncDecAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.EncDecAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.SelfAttention),
  ({"num_heads": 16},
   @mesh_tensorflow.transformer.transformer_layers.EncDecAttention),
  ({"hidden_size": 2560},
   @mesh_tensorflow.transformer.transformer_layers.DenseReluDense),
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
  @mesh_tensorflow.transformer.transformer_layers.SelfAttention,
  @mesh_tensorflow.transformer.transformer_layers.EncDecAttention,
  @mesh_tensorflow.transformer.transformer_layers.DenseReluDense,
]
